/*
Copyright (c) 2022 <jihaixuan> <<EMAIL>>
*/

#include "co.h"
#include <stdlib.h>
#include <stdio.h>
#include <stdlib.h>
#include <setjmp.h>
#include <time.h>
#define POOLSIZE 1024
#define STACK_SIZE 1024 * 1024 * 64
#define ALLIGN_OF_16(ptr) ((void *)((((long int)ptr) & (~0xf)) - 0x8))
struct co
{
  const char *name;
  void *stack;
  int status;
  void (*func)(void *);
  void *arg;
  jmp_buf env;
};
static struct co main_co;
static struct co *current;
static struct co *pool[POOLSIZE] = {NULL};
int pool_index = 0;

static void switch_stack(struct co *co)
{
  current = co;
  longjmp(co->env, 1);
}
static struct co *create_co;
void wrap_func(struct co *co)
{
  create_co = co;
  if (setjmp(co->env) == 0)
  {
    switch_stack(&main_co);
  };
  co->func(co->arg);
  co->status = -1;
  switch_stack(&main_co);
}
struct co *co_start(const char *name, void (*func)(void *), void *arg)
{
  pool[0] = &main_co;
  if (setjmp(main_co.env) != 0)
    return create_co;
  struct co *c = malloc(sizeof(struct co));
  c->stack = malloc(STACK_SIZE);
  c->name = name;
  c->func = func;
  c->arg = arg;
  c->status = 1;
  void *align_stack = ALLIGN_OF_16(c->stack + STACK_SIZE);
  for (int i = 0; i < POOLSIZE; i++)
  {
    if (pool[i] == NULL)
    {
      pool[i] = c;
      break;
    }
  }
  current = c;
#if defined(__x86_64__) && !defined(__ILP32__)
  asm volatile(
      "movq %0,%%rsp\n"
      "movq %0,%%rbp\n"
      "movq %2,%%rdi\n"
      "jmp *%1\n"
      :
      : "a"(align_stack), "b"(wrap_func), "c"(c)
      :);
#elif defined(__i386__) || defined(__ILP32__)
  asm volatile(
      "movl %0,%%esp\n"
      "movl %0,%%ebp\n"
      "pushl %2\n"
      "call *%1\n"
      :
      : "a"(align_stack), "b"(wrap_func), "c"(c)
      :);
#endif
  return NULL;
}

void co_wait(struct co *co)
{
  while (co->status != -1)
  {
    co_yield();
  }
  free(co->stack);
  free(co);
  for (int i = 0; i < POOLSIZE; i++)
  {
    if (pool[i] == co)
    {
      pool[i] = NULL;
    }
  }
}

void co_yield()
{
  if (setjmp(current->env) != 0)
    return;
  for (int i = 0; i < POOLSIZE; i++)
  {
    pool_index++;
    pool_index %= POOLSIZE;
    if (pool[pool_index] != NULL && pool[pool_index]->status == 1)
    {
      switch_stack(pool[pool_index]);
    }
  }
}
